package com.example.yckjbigdataservice.utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import redis.clients.jedis.JedisCluster;
import redis.clients.jedis.params.SetParams;

import java.util.*;
import java.util.concurrent.*;

/**
 * 分布式锁
 */
public class ResilientRedisLock {
    private static final Logger LOG = LoggerFactory.getLogger(ResilientRedisLock.class);
    private final JedisCluster jedis;
    private final ThreadLocal<Map<String, Integer>> reentrantCount = ThreadLocal.withInitial(HashMap::new);
    private final Map<String, ScheduledFuture<?>> watchdogTasks = new ConcurrentHashMap<>();
    private final ScheduledExecutorService watchdogScheduler = Executors.newScheduledThreadPool(1);

    private static final Long RELEASE_SUCCESS = 1L;
    private static final String LOCK_SUCCESS = "OK";
    private static final String SET_IF_NOT_EXIST = "NX";
    private static final String SET_WITH_EXPIRE_TIME = "PX";

    // Lua 脚本：加锁（setnx + expire）
    /*private static final String LOCK_SCRIPT = """
        if redis.call('setnx', KEYS[1], ARGV[1]) == 1 then
            redis.call('pexpire', KEYS[1], ARGV[2])
            return 1
        else
            return 0
        end
        """;*/
    private static final String LOCK_SCRIPT =
            "if redis.call('setnx', KEYS[1], ARGV[1]) == 1 then " +
                    "redis.call('pexpire', KEYS[1], ARGV[2]); " +
                    "return 1; " +
                    "else return 0; end";

    // Lua 脚本：解锁（get+del）
    /*private static final String UNLOCK_SCRIPT = """
        if redis.call('get', KEYS[1]) == ARGV[1] then
            return redis.call('del', KEYS[1])
        else
            return 0
        end
        """;*/
    private static final String UNLOCK_SCRIPT =
            "if redis.call('get', KEYS[1]) == ARGV[1] then " +
                    "return redis.call('del', KEYS[1]); " +
                    "else return 0; end";

    public ResilientRedisLock(JedisCluster jedis) {
        this.jedis = jedis;
    }

    /**
     * 高容错加锁：可重入 + 自动续期 + 多次尝试获取
     */
    public boolean tryLock(String rawKey, String requestId, long expireMillis, long maxWaitMillis) {
        String key = wrapKey(rawKey);
        long start = System.currentTimeMillis();
        Map<String, Integer> lockMap = reentrantCount.get();

        if (lockMap.containsKey(key)) {
            //LOG.info("本地锁 Reentrant lock for key ： {} " , key);
            lockMap.put(key, lockMap.get(key) + 1);
            return true;
        }

        do {
            try {
                Object result = jedis.eval(LOCK_SCRIPT, Collections.singletonList(key),
                        Arrays.asList(requestId, String.valueOf(expireMillis)));
                //LOG.info("执行redis原子赋值的状态 = {} " , key);
                if ("1".equals(result.toString())) {
                    lockMap.put(key, 1);
                    startWatchdog(key, requestId, expireMillis);
                    return true;
                }
            } catch (Exception e) {
                e.printStackTrace();
                // 可能是 slot 迁移 or 瞬时网络问题
                try {
                    //Thread.sleep(100); // 退避
                    TimeUnit.MILLISECONDS.sleep(100);
                } catch (InterruptedException ex) {
                    throw new RuntimeException(ex);
                }
            }
        } while (System.currentTimeMillis() - start < maxWaitMillis);

        return false;
    }

    /**
     * 解锁：可重入 & 安全释放
     */
    public void releaseLock(String rawKey, String requestId) {
        String key = wrapKey(rawKey);
        Map<String, Integer> lockMap = reentrantCount.get();
        Integer count = lockMap.get(key);
        if (count == null) return;

        if (count > 1) {
            lockMap.put(key, count - 1);
            return;
        }

        Object result = jedis.eval(UNLOCK_SCRIPT, Collections.singletonList(key), Collections.singletonList(requestId));
        LOG.debug("Unlock result: {}", result.toString());
        stopWatchdog(key);
        lockMap.remove(key);
        if (lockMap.isEmpty()) {
            reentrantCount.remove();
        }
    }

    /**
     * 启动 Watchdog：周期性续租
     */
    private void startWatchdog(String key, String requestId, long expireMillis) {
        long interval = expireMillis * 2 / 3;

        ScheduledFuture<?> future = watchdogScheduler.scheduleAtFixedRate(() -> {
            try {
                String current = jedis.get(key);
                if (requestId.equals(current)) {
                    jedis.pexpire(key, expireMillis);
                }
            } catch (Exception e) {
                // 忽略 Watchdog 异常
            }
        }, interval, interval, TimeUnit.MILLISECONDS);

        watchdogTasks.put(key, future);
    }

    private void stopWatchdog(String key) {
        ScheduledFuture<?> task = watchdogTasks.remove(key);
        if (task != null) task.cancel(false);
    }

    private String wrapKey(String key) {
        // 强制绑定 Redis cluster hash slot，避免跨节点操作
        return "lock:{" + key + "}";
    }


    /**
     * 释放分布式锁
     * @param lockKey 锁
     * @param requestId 请求标识
     * @return 是否释放成功
     */
    public boolean releaseDistributedLock( String lockKey, String requestId) {

        String script = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end";
        Object result = jedis.eval(script, Collections.singletonList(lockKey), Collections.singletonList(requestId));

        if (RELEASE_SUCCESS.equals(result)) {
            return true;
        }
        return false;

    }

    /**
     * 尝试获取分布式锁
     * @param lockKey 锁
     * @param requestId 请求标识
     * @param expireTime 超期时间
     * @return 是否获取成功
     */
    public boolean tryGetDistributedLock( String lockKey, String requestId, int expireTime) {
        SetParams params = new SetParams();
        params.ex(expireTime); // 设置超时时间
        params.nx(); // 若锁不存在才进行写操作
        params.px(expireTime);
        String result = jedis.set(lockKey, requestId, params);
        return LOCK_SUCCESS.equals(result);

    }

}

